import os
import shutil
from ultralytics import YOLO

# 加载模型
try:
    model = YOLO('best.pt')
except FileNotFoundError:
    print("未找到 best.pt 模型文件，请检查文件路径。")
    exit(1)

# 定义图片文件夹和目标文件夹
image_folder = 'Data/Photo'
target_folder_less_1 = 'Fail'
target_folder_more_1 = 'Success'

# 创建目标文件夹（如果不存在）
for folder in [target_folder_less_1, target_folder_more_1]:
    if not os.path.exists(folder):
        os.makedirs(folder)

# 遍历图片文件夹中的所有图片
for filename in os.listdir(image_folder):
    if filename.endswith(('.jpg', '.jpeg', '.png')):
        image_path = os.path.join(image_folder, filename)

        try:
            # 进行目标检测
            results = model(image_path)

            # 获取检测到的目标数
            for result in results:
                num_detections = len(result.boxes)

                # 如果目标数小于 1，复制图片到目标文件夹（少于 1 个目标的文件夹）
                if num_detections < 1:
                    shutil.copy2(image_path, os.path.join(target_folder_less_1, filename))
                    print(f"Copied {filename} to {target_folder_less_1}")
                # 如果目标数大于 1，复制图片到另一个目标文件夹（多于 1 个目标的文件夹）
                elif num_detections > 1:
                    shutil.copy2(image_path, os.path.join(target_folder_more_1, filename))
                    print(f"Copied {filename} to {target_folder_more_1}")
        except Exception as e:
            print(f"处理图片 {filename} 时出错: {e}")